
from datetime import datetime, timedelta

import pandas as pd


def split_weeks(df):
    def week_to_date(year, week):
        first_day_of_year = datetime(year, 1, 1)
        if first_day_of_year.weekday() <= 3:
            first_day_of_year -= timedelta(days=first_day_of_year.weekday())
        else:
            first_day_of_year += timedelta(days=7 - first_day_of_year.weekday())
        first_day_of_week = first_day_of_year + timedelta(weeks=int(week) - 1)
        return first_day_of_week

    # Добавляем дату начала недели
    df['start_date'] = df.apply(lambda row: week_to_date(row['year'], row['week']), axis=1)

    # Функция для расчета дней в месяце для каждой недели
    def days_in_month(start_date):
        end_date = start_date + timedelta(days=6)  # Конец недели
        current_date = start_date
        rows = []

        # Проверяем, является ли это последней неделей года
        if current_date.year != (current_date + timedelta(days=6)).year:
            # Вся неделя считается за текущий год
            rows.append((current_date.month, 7, current_date.year))
        else:
            while current_date <= end_date:
                month = current_date.month
                year = current_date.year
                month_days = 0

                # Считаем дни до конца месяца или недели
                while current_date <= end_date and current_date.month == month:
                    month_days += 1
                    current_date += timedelta(days=1)

                rows.append((month, month_days, year))

        return rows

    # Применяем функцию и расширяем датафрейм
    temp_df = df['start_date'].apply(days_in_month)
    expanded_data = []

    for i, row in enumerate(temp_df):
        for month, days, year in row:
            expanded_data.append([year, df.loc[i, 'week'], df.loc[i, 'start_date'], month, days])

    # Создаем новый датафрейм из расширенных данных
    expanded_df = pd.DataFrame(expanded_data, columns=['year', 'week', 'start_date', 'month', 'days_in_month'])
    return expanded_df

def split_qtty_to_weeks(df,service):
    prev = None
    for ind, row in df.iterrows():
        week_group = df[(df['year'] == row['year']) & (df['week'] == row['week'])]
        if week_group.shape[0] > 1:
            total_services = int(week_group[service].max())
            if not prev:
                prev = round(total_services * row['days_in_month'] / 7)
                df.at[ind, service] = prev
            else:
                # Последней строке присваиваем оставшееся количество услуг
                df.at[ind, service] = total_services - prev
                prev = None
    return df

def split_qtty_to_weeks(df,service):
    prev = None
    for ind, row in df.iterrows():
        week_group = df[(df['year'] == row['year']) & (df['week'] == row['week'])]
        if week_group.shape[0] > 1:
            total_services = int(week_group[service].max())
            if not prev:
                prev = round(total_services * row['days_in_month'] / 7)
                df.at[ind, service] = prev
            else:
                # Последней строке присваиваем оставшееся количество услуг
                df.at[ind, service] = total_services - prev
                prev = None
    return df

def split_predicted_qtty_to_weeks(df):
    prev = None
    for ind, row in df.iterrows():
        week_group = df[(df['year'] == row['year']) & (df['week'] == row['week'])]
        if week_group.shape[0] > 1:
            total_services = int(week_group['qtty'].max())
            if not prev:
                prev = round(total_services * row['days_in_month'] / 7)
                df.at[ind,'qtty'] = prev
            else:
                # Последней строке присваиваем оставшееся количество услуг
                df.at[ind,'qtty'] = total_services - prev
                prev = None
    return df